import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error
from scipy.optimize import minimize

# 读取输入数据
data_df = pd.read_csv("./data/source.csv")
inputs = data_df[["集水面积", "Sr", "Ks", "蒸发量", "降雨量(24h)"]].values
true_outputs = data_df["洪峰流量"].values

# 读取参数上下限
limit_df = pd.read_csv("./data/limit.csv", index_col=0)
lower_limit = limit_df.iloc[0].values
upper_limit = limit_df.iloc[1].values

# 参数名称
parameter_names = [
    "B",
    "alpha0",
    "JL",
    "KKS",
    "KKSS",
    "CS",
    "z1",
    "z2",
    "z3",
    "ni1",
    "ni2",
    "ni3",
    "H01",
    "H02",
    "H03",
    "Hc1",
    "Hc2",
    "Hc3",
    "Hr1",
    "Hr2",
    "Hr3",
    "delta1",
    "delta2",
    "delta3",
    "LS",
    "LSS",
    "L",
    "Gam1",
    "Gam2",
    "Gam3",
    "EC",
    "KKG",
    "LG",
]


# 定义模型函数：计算洪峰流量
def model(params, inputs):
    # 将33个参数与输入变量结合，计算洪峰流量（这里我们可以使用一个简单的线性模型进行示范）
    assert len(params) == 33, "参数的数量必须为33个"

    # 在实际情况中，模型会根据参数来计算洪峰流量
    # 这里只是一个示例，将参数和输入值相乘并求和来模拟洪峰流量的计算
    param_weights = params[:5]  # 假设前5个参数与输入特征相关
    output = np.dot(inputs, param_weights)  # 使用输入特征与权重相乘
    return output


# 损失函数：计算MSE
def loss_function(params, inputs, true_outputs):
    predictions = model(params, inputs)
    mse = mean_squared_error(true_outputs, predictions)
    return mse


# 参数初始化
initial_params = np.random.uniform(low=lower_limit, high=upper_limit, size=33)

# 使用梯度下降法优化参数
result = minimize(
    loss_function,
    initial_params,
    args=(inputs, true_outputs),
    bounds=list(zip(lower_limit, upper_limit)),
    method="L-BFGS-B",
)

# 获取优化后的参数
optimized_params = result.x

# 计算优化后的洪峰流量预测
optimized_predictions = model(optimized_params, inputs)

# 计算优化后的MSE
optimized_mse = mean_squared_error(true_outputs, optimized_predictions)
print(f"Optimized MSE: {optimized_mse:.4f}")

# 保存优化后的参数
optimized_param_df = pd.DataFrame(
    {"Parameter": parameter_names, "Optimized Value": optimized_params}
)

# 保存为CSV文件
output_filename = "./output/optimized_parameters.csv"
optimized_param_df.to_csv(output_filename, index=False)
print(f"Optimized parameters saved to {output_filename}")

# 可视化MSE下降过程
# 绘制损失函数随迭代的变化（模拟梯度下降过程中的损失变化）
loss_values = []

# 使用梯度下降算法进行模拟
for i in range(100):  # 模拟100次梯度下降更新
    result = minimize(
        loss_function,
        initial_params,
        args=(inputs, true_outputs),
        bounds=list(zip(lower_limit, upper_limit)),
        method="L-BFGS-B",
    )
    loss_values.append(result.fun)  # 保存每次的MSE
    initial_params = result.x  # 更新参数

plt.plot(loss_values)
plt.title("Loss Function (MSE) Over Iterations")
plt.xlabel("Iterations")
plt.ylabel("MSE")
plt.grid(True)
plt.show()

# 保存模型公式（简单的线性模型公式）
model_formula = "洪峰流量 = " + " + ".join(
    [
        f"{param:.4f} * {name}"
        for param, name in zip(optimized_params[:5], parameter_names[:5])
    ]
)
print(f"Model Formula: {model_formula}")

# 保存模型公式
with open("./output/model_formula.txt", "w") as f:
    f.write(model_formula)
    f.write("\n")
    f.write(f"Optimized MSE: {optimized_mse:.4f}")
